# -*- coding: utf-8 -*-
"""
Created on Thu Oct 01 05:14:12 2015

@author: zhonc0a
"""
import matplotlib.pyplot as plt
import numpy as np
import scipy.linalg as lina
import math
import collections
import os
import shutil
a = [0.3557421, 0.0, 0.0]
b = [0.1469838, 0.4418496, 0.0]
a1 = [0.3679705, -0.0027175, 0.0]
b1 = [-0.0121104, 0.4557614, 0.0]
a2 = [0.2717125, 0.0, 0.0]
b2 = [0.0, 0.4706183, 0.0]
a3 = [0.3223121, 0.0, 0.0]
b3 = [0.1611561, 0.2791305, 0]


class Brillouin:

    def __init__(self):

        self.num_layer = 4
        self.num_kpt_a = 10
        self.layer = []
        self.layer_frac = []
        self.radi_kpt = []
        self.radi_kpt_frac = []
        self.peri_kpt = []
        self.peri_kpt_frac = []
        self.layer_kpt = []
        self.layer_kpt_frac = []
        # find the symmetry
        #    calculate the angle between two vector
    def read_cell_param(self, filename):
        switch = 0
        with open(filename, 'r') as f:
            for line in f:
                if 'reciprocal lattice vectors' in line:
                    switch = 2
                    continue
                if switch == 2:
                    rv_a = map(float, line.strip().split()[3:])
                    switch -= 1
                    continue
                if switch == 1:
                    rv_b = map(float, line.strip().split()[3:])
                    switch -= 1
                    continue
        self.rv_a = np.array(rv_a[:2])
        self.rv_b = np.array(rv_b[:2])
    
    # generate mesh based on cell parameter and num_layer, num_kpt_a    
    def mesh_gen(self):
        rv_a = self.rv_a
        rv_b = self.rv_b
        self.rv_c = np.array([0, 0])       
        def angle(pt1, pt2):
            x1, y1 = pt1
            x2, y2 = pt2
            inner_product = x1 * x2 + y1 * y2
            len1 = math.hypot(x1, y1)
            len2 = math.hypot(x2, y2)
            return math.acos(inner_product / (len1 * len2))

        # find the intersection of two line that perpendicular and
        # intersect with vector a and b
        def intersect(a, b):
            param = np.array([[a[0], a[1]], [b[0], b[1]]])
            right = np.array([a[0]**2 + a[1]**2, b[0]**2 + b[1]**2])
            try:
                result = lina.solve(param, right)
            except lina.LinAlgError:
                return [None, None]
            return result
        angle_ab = angle(self.rv_a, self.rv_b)
        self.angle_ab = angle_ab
        # add the third axis, make axis a b c always counter-clockwise,
        # if angle between a and b > 90, then make b as c.
        # In this case, rv_b will different from self.rv_b
        if round(angle_ab, 3) < round(math.pi / 2, 3):
            rv_c = rv_b - rv_a
            self.rv_c = rv_c
        elif round(angle_ab, 3) > round(math.pi / 2, 3):
            rv_c = rv_b
            rv_b = rv_c + rv_a
            self.rv_c = rv_b
        if round(angle_ab, 3) == round(math.pi / 2, 3):
            self.sym = 'rec'
        elif (round(angle_ab, 3) == round(math.pi / 3, 3) or
              round(angle_ab, 3) == round(math.pi * 2 / 3, 3)) and  \
              round(math.hypot(*self.rv_a), 3) == round(math.hypot(*self.rv_b), 3):
            self.sym = 'hex'
        else:
            self.sym = 'other'
        # find the lines that defined the brillouin zone:
        if self.sym == 'rec':
            brillouin_corner = [intersect(rv_a/2, rv_b/2),
                                intersect(rv_a/-2, rv_b/2),
                                intersect(rv_a/-2, rv_b/-2),
                                intersect(rv_a/2, rv_b/-2)]
            self.radi = [rv_a/2, brillouin_corner[0], rv_b/2]
            self.peri = [rv_a/2, brillouin_corner[0], rv_b/2]
        if self.sym == 'hex':
            brillouin_corner = [intersect(rv_a/2, rv_b/2),
                                intersect(rv_b/2, rv_c/2),
                                intersect(rv_c/2, rv_a/-2),
                                intersect(rv_a/-2, rv_b/-2),
                                intersect(rv_b/-2, rv_c/-2),
                                intersect(rv_c/-2, rv_a/2)]
            self.radi = [rv_a/2, brillouin_corner[0]]
            self.peri = [rv_a/2, brillouin_corner[0]]
        if self.sym == 'other':
            brillouin_corner = [intersect(rv_a/2, rv_b/2),
                                intersect(rv_b/2, rv_c/2),
                                intersect(rv_c/2, rv_a/-2),
                                intersect(rv_a/-2, rv_b/-2),
                                intersect(rv_b/-2, rv_c/-2),
                                intersect(rv_c/-2, rv_a/2)]
            self.radi = [rv_a/2, brillouin_corner[0], rv_b/2,
                           brillouin_corner[1], rv_c/2, brillouin_corner[2]]
            self.peri = [rv_a/2, brillouin_corner[0], brillouin_corner[1],
                         brillouin_corner[2], rv_a/-2]
        self.brillouin_corner = brillouin_corner
        for i in range(self.num_layer):
            fract = ((i+1.0)/(self.num_layer+1.0))
            layer = [node*fract for node in self.peri]
            self.layer.append(layer)
        #g enerate seperate kpt and calculate num of kpt on each path:
        self.spacing = math.hypot(*(self.rv_a/2))/self.num_kpt_a

        def kpt_gen(endpt, startpt=np.array([0, 0])):
            endpt = np.array(endpt)
            startpt = np.array(startpt)
            num_kpt = int(round(math.hypot(*(endpt-startpt))/self.spacing, 0))
            actual_spc = (endpt-startpt) / num_kpt
            kpt = []
            for i in range(num_kpt+1):
                kpt.append(endpt-actual_spc*i)
            return kpt

        for kpt in self.radi:
            self.radi_kpt.append(kpt_gen(kpt))

        for i in range(len(self.peri)-1):
            self.peri_kpt.append(kpt_gen(self.peri[i+1], self.peri[i])[::-1])

        for layer in self.layer:
            layer_kpt = []
            for i in range(len(layer)-1):
                layer_kpt += layer_kpt+kpt_gen(layer[i+1], layer[i])[::-1]
            layer_kpt = [self.sig_fig(point) for point in layer_kpt]
            layer_kpt = list(collections.OrderedDict.fromkeys(layer_kpt))
            layer_kpt = [np.array(pt) for pt in layer_kpt]
            self.layer_kpt.append(layer_kpt)

    def sig_fig(self, point):
        point[0] = round(point[0], 5)
        point[1] = round(point[1], 5)
        return tuple(point)
        
    # make data points fill all the brillouin zone according to their sym
    # the input data should be a 2D-list, and the first two elements of
    # each item in list are coordinates
    def expand_bnz(self, data):
        def rotate(point, angle):
            new_point = list(point)
            coord = np.array(point[:2])
            rot_mat = np.array([[math.cos(angle), -math.sin(angle)],
                               [math.sin(angle), math.cos(angle)]])
            new_coord = np.dot(rot_mat, coord)
            new_point[:2] = new_coord
            return new_point

        def reflex(point, vect):
            new_point = list(point)
            coord = np.array(point[:2])
            new_coord = 2*np.dot(coord, vect)/np.dot(vect, vect)*vect - coord
            new_point[:2] = new_coord
            return new_point
        if self.sym == 'other':
            data_rot = [rotate(point, math.pi) for point in data]
            new_data = data+data_rot
        if self.sym == 'rec':
            refl_ax = self.rv_b
            data_mirror = [reflex(point, refl_ax) for point in data]
            new_data = data+data_mirror[::-1]
            data_rot = [rotate(point, math.pi) for point in new_data]
            new_data += data_rot
        if self.sym == 'hex':
            refl_ax = self.rv_a+self.rv_b
            angle = math.pi/3
            data_mirror = [reflex(point, refl_ax) for point in data]
            new_data = data+data_mirror[::-1]
            tmp_data = list(new_data)
            for i in range(5):
                data_rot = [rotate(point, angle) for point in tmp_data]
                tmp_data = data_rot
                new_data = new_data + data_rot

        # try to remove the duplicates generated by sym operation.
        new_data = [self.sig_fig(point) for point in new_data]
        new_data = list(collections.OrderedDict.fromkeys(new_data))
        return new_data + new_data[:1]

    def brillouin_plot(self):
        brillouin_path = zip(
            *(self.brillouin_corner + [self.brillouin_corner[0]]))
        # plot the mesh layer
        full_layer = [self.expand_bnz(path) for path in self.layer]
        layer_path = [zip(*path) for path in full_layer]
        for path in layer_path:
            plt.plot(*path)
        # plot the radi points
        full_radi_kpt = [self.expand_bnz(path) for path in self.radi_kpt]
        radi_path = [zip(*path) for path in full_radi_kpt]
        for path in radi_path:
            plt.scatter(*path, s=1)
        # plot the peri points
        full_peri_kpt = [self.expand_bnz(path) for path in self.peri_kpt]
        peri_path = [zip(*path) for path in full_peri_kpt]
        for path in peri_path:
            plt.scatter(*path, s=1)
        # plot the layer points
        full_layer_kpt = [self.expand_bnz(pt) for pt in self.layer_kpt]
        layer_kpt = [zip(*pt) for pt in full_layer_kpt]
        for kpt in layer_kpt:
            plt.scatter(*kpt, s=1)
        plt.arrow(0, 0, *self.rv_a, color='black', linewidth=1)
        plt.arrow(0, 0, *self.rv_b, color='black', linewidth=1)
        plt.arrow(0, 0, *self.rv_c, color='black', linewidth=1)
        
        plt.annotate('a*', self.rv_a, self.rv_a * 1.1,
                     style='italic', fontsize='large')
        plt.annotate('b*', self.rv_b, self.rv_b * 1.1,
                     style='italic', fontsize='large')
        if round(self.angle_ab, 3) < round(math.pi / 2, 3):
            plt.annotate('a* - b*', self.rv_c, self.rv_c * 1.1,
                         style='italic', fontsize='large')
        elif round(self.angle_ab, 3) > round(math.pi / 2, 3):
            plt.annotate('a* + b*', self.rv_c, self.rv_c * 1.1,
                         style='italic', fontsize='large')
        xmax = max(self.rv_a[0], self.rv_b[0])
        ymax = max(self.rv_a[1], self.rv_b[1])
        plt.axis([-1.1*xmax, 1.1*xmax, -1.1*ymax, 1.1*ymax])
        plt.plot(*(brillouin_path + ['r-']))


        
    def cart2frac(self):
        coord_mat = np.array([self.rv_a,self.rv_b])
        self.radi_frac = np.dot(np.array(self.radi), lina.inv(coord_mat))
        self.peri_frac = np.dot(np.array(self.peri), lina.inv(coord_mat))
        for layer in self.layer:
            self.layer_frac.append(np.dot(layer, lina.inv(coord_mat)))
        for path in self.radi_kpt:
            self.radi_kpt_frac.append(np.dot(path, lina.inv(coord_mat)))
        for path in self.peri_kpt:
            self.peri_kpt_frac.append(np.dot(path, lina.inv(coord_mat)))
        for path in self.layer_kpt:
            self.layer_kpt_frac.append(np.dot(path, lina.inv(coord_mat)))
        # make a dictionary to collect all path
        pd = {'radi_{:02d}'.format(i): self.radi_frac[i]
              for i in range(len(self.radi_frac))}.copy()
        pd.update({'peri_{:02d}'.format(i): np.array([self.peri_frac[i], self.peri_frac[i+1]])
                   for i in range(len(self.peri_frac)-1)})
        pd.update({'layer_{:02d}'.format(i): self.layer_frac[i]
                   for i in range(len(self.layer_frac))})
        self.path_dict = pd
        kd = {'radi_{:02d}'.format(i): self.radi_kpt_frac[i]
              for i in range(len(self.radi_kpt_frac))}.copy()
        kd.update({'peri_{:02d}'.format(i): self.peri_kpt_frac[i]
                   for i in range(len(self.peri_kpt_frac))})
        kd.update({'layer_{:02d}'.format(i): self.layer_kpt_frac[i]
                   for i in range(len(self.layer_kpt_frac))})
        self.kpt_dict = kd
    # 
    def gen_VASP_inp(self):
        '''
        read IBZKPT file
        make directory for radi peri and layer (one directory for each layer)
        cd into each directory and write corresponding kpoint 
        modified the INCAR file and add the ICHARGE = 11 and NELMIN = 5 keyword
        copy the WAVECAR, POSCAR, POTCAR and CHGCAR file to each directory
        '''
        with open('IBZKPT', 'r') as f:
            ibzkpt = list(f)
        with open('INCAR', 'r') as f:
            incar = list(f)
        origin_num_kpt = int(ibzkpt[1])
        layer_dir = ['layer_{:02d}'.format(i) for i in range(len(self.layer_kpt))]
        radi_dir = ['radi_{:02d}'.format(i) for i in range(len(self.radi_kpt))]
        peri_dir = ['peri_{:02d}'.format(i) for i in range(len(self.peri_kpt))]
        all_dir = radi_dir+peri_dir+layer_dir
        print all_dir
        base_dir = os.getcwd()
        
        for dir in all_dir:
            dir_name = '2D_'+dir
            if not os.path.exists(dir_name):
                os.makedirs(dir_name)
            
            shutil.copy('WAVECAR', dir_name)
            shutil.copy('CHGCAR', dir_name)
            shutil.copy('POSCAR', dir_name)
            shutil.copy('POTCAR', dir_name)
            
            kpt_list = [list(i)+[0.0, 0] for i in self.kpt_dict[dir]]
            ibzkpt[1] = str(origin_num_kpt+len(kpt_list))+'\n'
            os.chdir(dir_name)
            kpt_file = open('KPOINTS', 'w')
            for line in ibzkpt:
                kpt_file.write(line)
            for kpt in kpt_list:
                kpt_file.write("{:>20.14f}{:>20.14f}{:>20.14f}{:14d}\n".format(*kpt))
            kpt_file.close()
            incar_file = open('INCAR', 'w')
            for line in incar:
                if 'ICHARGE' in line or 'NELMIN' in line:
                    line = ''
                incar_file.write(line)
            incar_file.write('ICHARGE=11\n')
            incar_file.write('NELMIN=5\n')
            
                    
            kpt_file.close()
            os.chdir(base_dir)
            
#            
#            os.chdir(dir)
#            open('KPOINTS', 'w')
            
            


#    intersect_ab=intersect(mid_a,mid_b)
br = Brillouin()
br.read_cell_param('OUTCAR')
br.mesh_gen()
br.cart2frac()
br.brillouin_plot()
br.gen_VASP_inp()
